import os
import sys
import shutil
import re

# 根据文件名中的关键字，自动分类文件

SRC_DIR = sys.argv[1]
DIST_DIR = "dist"

print("SRC_DIR", SRC_DIR)

KEYWORDS = ["工业投资运行预计", "5亿元以上项目", "能源外工业投资预计"]
# KEYWORDS = ["未开工", "已开工未入库", "投资完成滞后"]
TARGET_NAMES = ["gy", "51", "ne"]


def prepare_dir(dir: str, renew=True) -> None:
    """
    准备目标文件夹：
    - 存在的话
        - 如果要更新，先删除再新建
        - 如果不更新，不需要新建
    - 不存在的话，新建
    :param dir: 目标文件夹路径
    :param renew: 是否删除重新创建
    """
    if os.path.exists(dir):
        if renew:
            shutil.rmtree(dir)
            os.mkdir(dir)
    else:
        os.mkdir(dir)


def get_files(dir, ext, recursive=False):
    """
    获取指定目录下指定扩展名的文件列表，默认不递归
    :param dir: 目录路径
    :param ext: 扩展名，带点 `.xlsx`
    :param recursive: 是否递归查找
    :return: 文件列表
    """
    all_files = []

    if recursive:
        # 遍历指定目录及其子目录下的所有文件
        for root, _, files in os.walk(dir):
            for file in files:
                # 使用os.path.join拼接完整的文件路径
                all_files.append(os.path.join(root, file))
    else:
        # 只取一层文件夹的文件
        all_files = map(lambda x: os.path.join(dir, x), os.listdir(dir))

    # 排除扩展名不匹配的文件，还有 `~$` 开头的临时文件
    valid_ext_files = filter(
        lambda x: x.endswith(ext) and not os.path.basename(x).startswith(".~"),
        all_files,
    )

    return list(valid_ext_files)


def try_file_path(file_path):
    """
    测试是否存在文件，如果存在，返回带括号的文件名
    :param file_path: 文件路径
    :return: 新文件名
    """
    tmp_path = file_path
    dir_and_name, ext = os.path.splitext(tmp_path)
    i = 1
    while os.path.exists(tmp_path):
        tmp_path = f"{dir_and_name}({i}){ext}"
        i += 1
    return tmp_path


def check_filenames(files):
    paths = [f"{DIST_DIR}/{k}" for k in TARGET_NAMES]
    [os.mkdir(p) for p in paths if not os.path.exists(p)]

    for f in files:
        for k in KEYWORDS:
            filename = os.path.basename(f)
            mat = re.search(k, filename)
            if mat is not None:
                new_key = TARGET_NAMES[KEYWORDS.index(k)]
                print("Move", f, k)
                target = try_file_path(f"{DIST_DIR}/{new_key}/{filename}")
                shutil.copy(f, target)


def main():
    prepare_dir(DIST_DIR)

    files = get_files(SRC_DIR, ".xlsx", True)
    print(files)
    check_filenames(files)

    print("Done")


if __name__ == "__main__":
    main()
